package com.joker.airoleplayplatformback.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket连接管理工具类
 * 负责管理WebSocket连接、消息发送和会话状态
 *
 * @author feixiang.li
 * @since 2025-09-27
 */
@Slf4j
@Component
public class WebSocketConnectionManager {

    private final ObjectMapper objectMapper = new ObjectMapper();

    // 存储WebSocket会话
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    // 存储会话对应的用户信息
    private final Map<String, Long> sessionUserMap = new ConcurrentHashMap<>();

    // 存储会话最后心跳时间
    private final Map<String, Long> sessionHeartbeatMap = new ConcurrentHashMap<>();

    /**
     * 注册WebSocket会话
     */
    public void registerSession(String sessionId, WebSocketSession session) {
        sessions.put(sessionId, session);
        sessionHeartbeatMap.put(sessionId, System.currentTimeMillis());
        log.info("WebSocket会话已注册: {}", sessionId);
    }

    /**
     * 注销WebSocket会话
     */
    public void unregisterSession(String sessionId) {
        sessions.remove(sessionId);
        sessionUserMap.remove(sessionId);
        sessionHeartbeatMap.remove(sessionId);
        log.info("WebSocket会话已注销: {}", sessionId);
    }

    /**
     * 设置会话用户信息
     */
    public void setSessionUser(String sessionId, Long userId) {
        sessionUserMap.put(sessionId, userId);
    }

    /**
     * 获取会话用户信息
     */
    public Long getSessionUser(String sessionId) {
        return sessionUserMap.get(sessionId);
    }

    /**
     * 更新心跳时间
     */
    public void updateHeartbeat(String sessionId) {
        sessionHeartbeatMap.put(sessionId, System.currentTimeMillis());
    }

    /**
     * 获取心跳时间
     */
    public Long getHeartbeatTime(String sessionId) {
        return sessionHeartbeatMap.get(sessionId);
    }

    /**
     * 检查会话是否存在且连接正常
     */
    public boolean isSessionActive(String sessionId) {
        WebSocketSession session = sessions.get(sessionId);
        return session != null && session.isOpen();
    }

    /**
     * 发送文本消息到指定会话
     */
    public void sendMessage(String sessionId, Object message) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                String jsonMessage = objectMapper.writeValueAsString(message);
                session.sendMessage(new TextMessage(jsonMessage));
                log.debug("消息发送成功: sessionId={}, message={}", sessionId, message);
            } catch (IOException e) {
                log.error("发送消息失败: sessionId={}", sessionId, e);
            }
        } else {
            log.warn("会话不存在或已关闭: {}", sessionId);
        }
    }

    public void sendMessageByte(String sessionId, byte[] message) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new BinaryMessage(message));
                log.debug("消息发送成功: sessionId={}, message={}", sessionId, message);
            } catch (IOException e) {
                log.error("发送消息失败: sessionId={}", sessionId, e);
            }
        } else {
            log.warn("会话不存在或已关闭: {}", sessionId);
        }
    }


    /**
     * 发送二进制消息（音频数据）到指定会话
     */
    public void sendBinaryMessage(String sessionId, byte[] audioData) {
        WebSocketSession session = sessions.get(sessionId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new BinaryMessage(audioData));
                log.debug("音频数据发送成功: sessionId={}, size={}", sessionId, audioData.length);
            } catch (IOException e) {
                log.error("发送音频数据失败: sessionId={}", sessionId, e);
            }
        } else {
            log.warn("会话不存在或已关闭: {}", sessionId);
        }
    }

    /**
     * 发送错误消息到指定会话
     */
    public void sendErrorMessage(String sessionId, String errorMessage) {
        sendMessage(sessionId, Map.of("type", "error", "message", errorMessage));
    }

    /**
     * 发送连接成功消息
     */
    public void sendConnectionMessage(String sessionId) {
        sendMessage(sessionId, Map.of("type", "connection", "status", "connected", "sessionId", sessionId));
    }

    /**
     * 发送初始化成功消息
     */
    public void sendInitSuccessMessage(String sessionId) {
        sendMessage(sessionId, Map.of("type", "init_success", "message", "会话初始化成功"));
    }

    /**
     * 发送语音录制开始消息
     */
    public void sendVoiceRecordingStartMessage(String sessionId) {
        sendMessage(sessionId, Map.of("type", "voice_recording", "status", "started"));
    }

    /**
     * 发送语音处理中消息
     */
    public void sendVoiceProcessingMessage(String sessionId) {
        sendMessage(sessionId, Map.of("type", "voice_processing", "status", "processing"));
    }

    /**
     * 发送音频停止消息
     */
    public void sendAudioStoppedMessage(String sessionId) {
        sendMessage(sessionId, Map.of("type", "audio_stopped", "status", "stopped"));
    }

    /**
     * 发送心跳响应消息
     */
    public void sendHeartbeatResponse(String sessionId) {
        sendMessage(sessionId, Map.of("type", "heartbeat_response", "timestamp", System.currentTimeMillis()));
    }

    /**
     * 发送AI回复消息
     */
    public void sendAiReplyMessage(String sessionId, String messageId, String content, String audioUrl, Long duration) {
        sendMessage(sessionId, Map.of("type", "ai_reply", "messageId", messageId, "content", content, "audioUrl", audioUrl, "duration", duration));
    }

    /**
     * 获取所有会话（供心跳检测使用）
     */
    public Map<String, WebSocketSession> getAllSessions() {
        return sessions;
    }

    /**
     * 获取心跳时间映射（供心跳检测使用）
     */
    public Map<String, Long> getHeartbeatMap() {
        return sessionHeartbeatMap;
    }

    /**
     * 获取会话数量
     */
    public int getSessionCount() {
        return sessions.size();
    }

    /**
     * 清理无效会话
     */
    public void cleanupInactiveSessions() {
        sessions.entrySet().removeIf(entry -> {
            String sessionId = entry.getKey();
            WebSocketSession session = entry.getValue();
            if (!session.isOpen()) {
                sessionUserMap.remove(sessionId);
                sessionHeartbeatMap.remove(sessionId);
                log.info("清理无效会话: {}", sessionId);
                return true;
            }
            return false;
        });
    }

    /**
     * 发送二进制音频数据
     */
    public void sendAudioData(String sessionId, byte[] audioData) {
        try {
            WebSocketSession session = sessions.get(sessionId);
            if (session != null && session.isOpen()) {
                BinaryMessage binaryMessage = new BinaryMessage(audioData);
                session.sendMessage(binaryMessage);
                log.info("发送音频数据: sessionId={}, size={} bytes", sessionId, audioData.length);
            } else {
                log.warn("会话不存在或已关闭，无法发送音频数据: {}", sessionId);
            }
        } catch (IOException e) {
            log.error("发送音频数据失败: sessionId={}", sessionId, e);
        }
    }

    /**
     * 发送MP3音频数据消息
     */
    public void sendMp3AudioMessage(String sessionId, byte[] mp3Data) {
        try {
            sendMessageByte(sessionId, mp3Data);
            log.info("发送MP3音频消息: sessionId={}, size={} bytes", sessionId, mp3Data.length);
        } catch (Exception e) {
            log.error("发送MP3音频消息失败: sessionId={}", sessionId, e);
        }
    }

    /**
     * 发送语音处理完成消息
     */
    public void sendVoiceProcessedMessage(String sessionId) {
        try {
            Map<String, Object> message = Map.of(
                    "type", "voice_processed",
                    "timestamp", System.currentTimeMillis()
            );
            sendMessage(sessionId, message);
            log.info("发送语音处理完成消息: sessionId={}", sessionId);
        } catch (Exception e) {
            log.error("发送语音处理完成消息失败: sessionId={}", sessionId, e);
        }
    }
}
